# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
SETTLE Constraint algorithm
"""

from typing import Union, Tuple, List, Dict

import numpy as np
from numpy import ndarray
import torch
import torch.nn.functional as F
from torch import Tensor, nn

from . import Constraint
from ...system import Molecule
from ...function import get_tensor


class EinsumWrapper(nn.Module):
    """Implement particular Einsum operation

    Args:

        equation (str):  en equation represent the operation.

    Supported Platforms:

        ``Ascend`` ``GPU``

    """
    def __init__(self, equation: str):
        super().__init__()
        self.equation = equation

    def forward(self, xy: Tuple[Tensor, Tensor]) -> Tensor:
        """Calculation for Einsum operation"""

        result = None
        if self.equation == 'ijk,ijk->ij':
            ijk1, ijk2 = xy
            ij = torch.sum(ijk1*ijk2, dim=-1)
            result = ij
        elif self.equation == 'ijk,ijl->ikl':
            ijk, ijl = xy
            ijkl1 = ijk.unsqueeze(-1).expand(*ijk.shape, ijl.shape[-1])
            ijkl2 = ijl.unsqueeze(-2).expand(*ijl.shape[:-1], ijk.shape[-1], ijl.shape[-1])
            result = (ijkl1 * ijkl2).sum(dim=1)
        else:
            raise NotImplementedError("This equation is not implemented")
        return result


class SETTLE(Constraint):
    """
    SETTLE constraint controller.

    `SETTLE: An Analytical Version of the SHAKE and RATTLE Algorithm for Rigid Water Models.`
        by Shuichi Miyamoto and Peter A. Kollman.

    Args:
        system (Molecule): Simulation system.
        index (Union[Tensor, ndarray, List[int]]): Array of settle index of shape `(C, 3)` or `(B, C, 3)`,
            and the type is int. If `None` is given, the `settle_index` in `system` will be used. Default: None
        distance (Union[Tensor, ndarray, List[float]]): Array of settle distance of shape (C, 2) or `(B, C, 2)`,
            and the type is float. If `None` is given, the `settle_dis` in `system` will be used. Default: None

    Inputs:
        - **coordinate** (Tensor) - The coordinates of the system.
        - **velocity** (Tensor) - The velocity of the system.
        - **force** (Tensor) - The force of the system.
        - **energy** (Tensor) - The energy of the system.
        - **kinetics** (Tensor) - The kinetics of the system.
        - **virial** (Tensor) - The virial of the system. Default: None
        - **pbc_box** (Tensor) - PBC box of the system. Default: None
        - **step** (int) - The step of the system. Default: 0

    Return:
        - coordinate (Tensor), Tensor of shape (B, A, D). Data type is float.
        - velocity (Tensor), Tensor of shape (B, A, D). Data type is float.
        - force (Tensor), Tensor of shape (B, A, D). Data type is float.
        - energy (Tensor), Tensor of shape (B, 1). Data type is float.
        - virial (Tensor), Tensor of shape (B, D). Data type is float.
        - pbc_box (Tensor), Tensor of shape (B, D). Data type is float.

    Supported Platforms:
        ``Ascend`` ``GPU``
    """

    def __init__(self,
                 system: Molecule,
                 index: Union[Tensor, ndarray, List[int]] = None, # (B, C, 3)
                 distance: Union[Tensor, ndarray, List[float]] = None, # (B, C, 2)
                 ):
        super().__init__(system, bonds='all-bonds')
        print('[MindSPONGE] The settle constraint is used for the molecule system.')

        # pylint: disable=invalid-name
        bsize = system.num_walker
        def _check_shape(value: Tensor, name: str, dim: int) -> Tensor:
            if value.ndim < 2 or value.ndim > 3:
                raise ValueError(f'The rank(ndim) of {name} should be 2 or 3 but got {value.ndim}')
            if value.ndim == 2:
                value = F.expand_dims(value, 0)
            if value.shape[0] != 1 and value.shape[0] != bsize:
                raise ValueError(f'The first dimension of {name} must be equal to 1 or '
                                 f'the number of multiple walkers of the simulation system ({bsize})'
                                 f'but got: {value.shape[0]}')
            if value.shape[-1] != dim:
                raise ValueError(f'The last dimension of {name} must be {dim} but got {value.shape[-1]}')
            return value

        index = get_tensor(index, torch.int32)
        if index is None:
            index = system.settle_index
        index = _check_shape(index, 'index', 3)

        distance = get_tensor(distance, torch.float32)
        if distance is None:
            distance = system.settle_length
        distance = _check_shape(distance, 'distance', 2)

        num_settle = index.shape[1]

        if distance.shape[-2] != num_settle:
            raise ValueError(f'The number of constraint in `distance` ({distance.shape[-2]}) does not match '
                             f'the number of constraint in `index` ({num_settle}).')

        self.num_constraints = num_settle * 3

        # (B, C)
        dis_legs = distance[..., 0]
        dis_base = distance[..., 1]

        # (B, C * 3)
        self.settle_index = index.reshape((bsize, -1))
        self.bs_index = F.expand_dims(self.settle_index[..., None], 2).expand((bsize, self.settle_index.shape[1], 3))
        self._mass = torch.take_along_axis(system.atom_mass,
                                          self.settle_index,
                                          axis=-1).reshape((bsize, num_settle, 3, 1))

        # (B, C, 3)
        self.target_yvector = torch.tensor([[[0, 1, 0]] * num_settle] * bsize, dtype=torch.float32)
        self.target_zvector = torch.tensor([[[0, 0, 1]] * num_settle] * bsize, dtype=torch.float32)

        # (B, 1, 4)
        self.identity_quaternion = torch.tensor([[[1, 0, 0, 0]]] * bsize, dtype=torch.float32)
        self.inverse_quaternion = torch.tensor([[[1, -1, -1, -1]]] * bsize, dtype=torch.float32)

        hx = torch.tensor(0.5 * dis_base.astype(np.float32), dtype=torch.float32)
        hy = torch.tensor(np.sqrt(dis_legs ** 2 - (0.5 * dis_base) ** 2).astype(np.float32), dtype=torch.float32)

        # (B, C, 1)
        self.mass_vertex = self._mass[..., 0, :]
        self.mass_base = self._mass[..., 1, :]

        # (B, C)
        self.rb = hy / (1 + 2 * self.mass_base[..., 0] / self.mass_vertex[..., 0])
        self.ra = hy - self.rb
        self.rc = hx

        # (B, C, 3, 3)
        self.crd_1 = torch.zeros((bsize, num_settle, 3, 3), dtype=torch.float32)
        self.crd_1[..., 0, 1] = self.ra
        self.crd_1[..., 1, 1] = -self.rb
        self.crd_1[..., 2, 1] = -self.rb
        self.crd_1[..., 1, 0] = -self.rc
        self.crd_1[..., 2, 0] = self.rc

        self.aindices = torch.tensor([0, 1], dtype=torch.int32)
        self.bindices = torch.tensor([0, 2], dtype=torch.int32)
        self.cindices = torch.tensor([0, 3], dtype=torch.int32)
        self.hindices = torch.tensor([1, 2], dtype=torch.int32)

        self.dot = EinsumWrapper('ijk,ijk->ij')
        self.einsum = EinsumWrapper('ijk,ijl->ikl')
        self.scatter_update = torch.scatter_


    def get_vector_transform(self, vec1, vec2):
        """ Get the transform quaternion from a vector to another.
        Args:
            vec1 (Tensor): The initial vector.
            vec2 (Tensor): The target vector.
        Returns:
            q (Tensor): The transform quaternion.
        """
        # (B, G, D)
        cross_vector = torch.cross(vec1, vec2, dim=-1)
        # (B, G, 1)
        cross_norm = torch.norm(cross_vector, dim=-1, keepdim=True)
        zero_index = torch.where(cross_norm == 0, torch.ones_like(cross_norm), torch.zeros_like(cross_norm))
        qs = torch.zeros((vec1.shape[0], vec1.shape[0], 1))
        # (B, G, D)
        qv = torch.zeros_like(vec1)
        # (B, G, 1)
        qs += torch.norm(vec1, dim=-1, keepdim=True) * torch.norm(vec2, dim=-1, keepdim=True)
        qs += self.dot((vec1, vec2))[..., None]
        # (B, G, D)
        qv += cross_vector
        # (B, G, 4)
        q = torch.cat((qs, qv), dim=-1)
        q = q * (1 - zero_index) + self.identity_quaternion * zero_index
        q /= torch.norm(q, dim=-1, keepdim=True)
        return q

    def get_inverse(self, quater):
        """ Get the inverse operation of a given quaternion.
        Args:
            quater (Tensor): The given quaternion.
        Returns:
            quater^{-1}
        """
        factor = torch.norm(quater, dim=-1, keepdim=True) ** 2
        return quater * self.inverse_quaternion / factor

    def quaternion_multiply(self, tensor_1, tensor_2):
        """ Calculate the quaternion multiplication.
        Args:
            tensor_1 (Tensor): The first quaternion, if the size of last dimension is 3, it will be padded to 4 auto.
            tensor_2 (Tensor): The second quaternion, if the size of last dimension is 3, it will be padded to 4 auto.
        Returns:
            q (Tensor): The quaternion product of tensor_1 and tensor_2.
        """
        # (B, G, 4)
        if tensor_1.shape[-1] == 3:
            tensor_1 = F.pad(tensor_1, (0, 0, 0, 0, 1, 0), mode='constant', value=0)
        if tensor_2.shape[-1] == 3:
            tensor_2 = F.pad(tensor_2, (0, 0, 0, 0, 1, 0), mode='constant', value=0)
        # (B, G, 1)
        s_1 = tensor_1[..., [0]]
        s_2 = tensor_2[..., [0]]
        # (B, G, 3)
        v_1 = tensor_1[..., 1:]
        v_2 = tensor_2[..., 1:]
        # (B, G, 1)
        s = s_1 * s_2
        # (B, G, 1)
        d = self.dot((v_1, v_2))[..., None]
        s -= d
        # (B, G, 3)
        v = torch.zeros_like(v_1)
        v += s_1 * v_2
        v += v_1 * s_2
        v += torch.cross(v_1, v_2, dim=-1)
        q = torch.cat((s, v), dim=-1)
        return q

    def hamiltonian_product(self, q, v):
        """ Perform hamiltonian product.
        Args:
            q (Tensor): The transform quaternion.
            v (Tensor): The vector to be transformed.
        Returns:
            qvq^{-1}
        """
        # (B, G, 4)
        iq = self.get_inverse(q)
        op1 = self.quaternion_multiply(v, iq)
        res = self.quaternion_multiply(q, op1)
        return res

    def group_hamiltonian_product(self, q, vec):
        """ Perform hamiltonian product in a 4-dimensional vector.
        Args:
            q (Tensor): The transform quaternion.
            vec (Tensor): The vector to be transformed.
        Returns:
            qvq^{-1}
        """
        # (B, G, 1, 4)
        group_q = q[..., None, :]
        # (B, G, a, 4)
        pad_q = F.pad(group_q, (0, 0, 0, 0, 0, 2, 0, 0), mode='circular')
        # (B, G*a, 4)
        res_q = pad_q.reshape((q.shape[0], q.shape[1]*3, q.shape[2]))
        iq = self.get_inverse(res_q)
        op1 = self.quaternion_multiply(vec.reshape((q.shape[0], q.shape[1]*3, vec.shape[2])), iq)
        # (B, G, a, D)
        res = self.quaternion_multiply(res_q, op1)[..., 1:].reshape(vec.shape)
        return res


    def get_transform(self, crd_):
        """ Get the transform between A0B0C0 and a0b0c0.
        Args:
            crd_ (Tensor): The given water molecular coordinates.
        Returns:
            This function will return the transform and inverse transform from crd_ SETTLE axis.
        """
        # (B, G, A, D)
        crd = crd_.reshape((crd_.shape[0], -1, 3, crd_.shape[-1]))
        # (B, G, D)
        mass_center = torch.sum(crd * self._mass, axis=-2) / self._mass.sum(axis=-2)
        # (B, G, D)
        cross_vector_1 = crd[..., 1, :] - crd[..., 0, :]
        cross_vector_2 = crd[..., 2, :] - crd[..., 1, :]
        cross_vector = torch.cross(cross_vector_1, cross_vector_2, dim=-1)
        cross_vector /= torch.norm(cross_vector, dim=-1, keepdim=True)
        # (B, G, 4)
        transform_1 = self.get_vector_transform(cross_vector, self.target_zvector)

        inverse_transform_1 = self.get_inverse(transform_1)
        # (B, G, D)
        oa_vector = crd[..., 0, :] - mass_center
        oa_vector /= torch.norm(oa_vector, dim=-1, keepdim=True)
        transformed_oa_vector = self.hamiltonian_product(transform_1, oa_vector)[..., 1:]
        transformed_oa_vector /= torch.norm(transformed_oa_vector, dim=-1, keepdim=True)
        # (B, G, 4)
        transform_2 = self.get_vector_transform(transformed_oa_vector, self.target_yvector)
        inverse_transform_2 = self.get_inverse(transform_2)
        # (B, G, 4)
        combine_transform = self.quaternion_multiply(transform_2, transform_1)
        combine_itransform = self.quaternion_multiply(inverse_transform_1, inverse_transform_2)
        return combine_transform, combine_itransform

    def apply_transform(self, q, vec):
        """ Apply the quaternion transform q to a given vector vec. """
        # (B, G, a, D)
        crd = vec.reshape((vec.shape[0], -1, 3, vec.shape[-1]))
        crd -= self.get_mass_center(crd)
        # (B, G, a, 1)
        norm_crd = torch.norm(crd, dim=-1, keepdim=True)
        # (B, G, a, D)
        crd /= norm_crd
        transform_crd = self.group_hamiltonian_product(q, crd)
        res = transform_crd * norm_crd
        # (B, A, D)
        return res.reshape(vec.shape)

    def _rotation(self, crd_1_0_, crd_1_1_):
        """ Get the rotation quaternion between triangles. """
        crd_1_0 = crd_1_0_.reshape((crd_1_0_.shape[0], -1, 3, crd_1_0_.shape[-1]))
        crd_1_1 = crd_1_1_.reshape((crd_1_1_.shape[0], -1, 3, crd_1_1_.shape[-1]))
        # (B, G)
        sin_phi = crd_1_1[..., 0, 2] / self.ra
        sin_phi = sin_phi.clamp(-0.999999, 0.999999)
        phi = torch.asin(sin_phi)
        sin_psi = (crd_1_1[..., 1, 2] - crd_1_1[..., 2, 2]) / 2 / self.rc / torch.sqrt(1 - sin_phi ** 2)
        sin_psi = sin_psi.clamp(-0.999999, 0.999999)
        psi = torch.asin(sin_psi)
        xb2 = -self.rc * torch.sqrt(1 - sin_psi ** 2)
        yb2 = -self.rb * torch.sqrt(1 - sin_phi ** 2) - self.rc * sin_phi * sin_psi
        yc2 = -self.rb * torch.sqrt(1 - sin_phi ** 2) + self.rc * sin_phi * sin_psi
        alpha = xb2 * (crd_1_0[..., 1, 0] - crd_1_0[..., 2, 0]) + \
                yb2 * (crd_1_0[..., 1, 1] - crd_1_0[..., 0, 1]) + \
                yc2 * (crd_1_0[..., 2, 1] - crd_1_0[..., 0, 1])
        beta = xb2 * (crd_1_0[..., 2, 1] - crd_1_0[..., 1, 1]) + \
               yb2 * (crd_1_0[..., 1, 0] - crd_1_0[..., 0, 0]) + \
               yc2 * (crd_1_0[..., 2, 0] - crd_1_0[..., 0, 0])
        gamma = (crd_1_0[..., 1, 0] - crd_1_0[..., 0, 0]) * crd_1_1[..., 1, 1] - \
                crd_1_1[..., 1, 0] * (crd_1_0[..., 1, 1] - crd_1_0[..., 0, 1]) + \
                (crd_1_0[..., 2, 0] - crd_1_0[..., 0, 0]) * crd_1_1[..., 2, 1] - \
                crd_1_1[..., 2, 0] * (crd_1_0[..., 2, 1] - crd_1_0[..., 0, 1])
        sin_theta = (alpha * gamma - beta * torch.sqrt(alpha ** 2 + beta ** 2 - gamma ** 2)) / (alpha ** 2 + beta ** 2)
        sin_theta = sin_theta.clamp(-0.999999, 0.999999)
        theta = torch.asin(sin_theta)

        # (B, G, 4)
        cos_phi = F.pad(torch.cos(phi[..., None] / 2), (0, 0, 0, 0, 0, 3), mode='constant', value=0)
        sin_phi = F.pad(torch.sin(phi[..., None] / 2), (0, 0, 0, 0, 0, 3), mode='constant', value=0)
        cos_psi = F.pad(torch.cos(psi[..., None] / 2), (0, 0, 0, 0, 0, 3), mode='constant', value=0)
        sin_psi = F.pad(torch.sin(psi[..., None] / 2), (0, 0, 0, 0, 0, 3), mode='constant', value=0)
        cos_theta = F.pad(torch.cos(theta[..., None] / 2), (0, 0, 0, 0, 0, 3), mode='constant',
                             value=0)
        sin_theta = F.pad(torch.sin(theta[..., None] / 2), (0, 0, 0, 0, 0, 3), mode='constant',
                             value=0)

        phi_minus_value = -sin_phi[..., self.aindices]
        psi_minus_value = -sin_psi[..., self.bindices]
        theta_minus_value = -sin_theta[..., self.cindices]
        phi_add_value = sin_phi[..., self.aindices[::-1]]
        psi_add_value = sin_psi[..., self.bindices[::-1]]
        theta_add_value = sin_theta[..., self.cindices[::-1]]

        sin_phi[..., self.aindices] += phi_minus_value
        sin_phi[..., self.aindices] += phi_add_value
        sin_psi[..., self.bindices] += psi_minus_value
        sin_psi[..., self.bindices] += psi_add_value
        sin_theta[..., self.cindices] += theta_minus_value
        sin_theta[..., self.cindices] += theta_add_value

        # (B, 4)
        quater_phi = cos_phi + sin_phi
        quater_psi = cos_psi + sin_psi
        quater_theta = cos_theta + sin_theta

        rotation = self.quaternion_multiply(quater_theta, self.quaternion_multiply(quater_phi, quater_psi))
        rotation /= torch.norm(rotation, dim=-1, keepdim=True)
        return rotation

    def _swap_h(self, crd_1_1_, crd_1_3_):
        """ Swap the H atoms if the resulted hydrogen from SETTLE are mixed. """
        # (B, G, a, D)
        #pylint: disable=invalid-name
        crd_1_1 = crd_1_1_.reshape((crd_1_1_.shape[0], -1, 3, crd_1_1_.shape[-1]))
        crd_1_3 = crd_1_3_.reshape((crd_1_3_.shape[0], -1, 3, crd_1_3_.shape[-1]))
        # (B, G)
        Hs_0_0 = crd_1_1[..., 1, 2]
        Hs_0_1 = crd_1_1[..., 2, 2]
        mask_0 = Hs_0_0 >= Hs_0_1
        Hs_1_0 = crd_1_3[..., 1, 2]
        Hs_1_1 = crd_1_3[..., 2, 2]
        mask_1 = Hs_1_0 >= Hs_1_1
        # (B, G, 1, 1)
        H_swap_mask = torch.where(mask_0 == mask_1, torch.zeros_like(mask_0), torch.ones_like(mask_0))[..., None, None]
        # (B, G, 2, D)
        h_minus_value = -crd_1_3[..., self.hindices, :] * H_swap_mask
        h_add_value = crd_1_3[..., self.hindices[::-1], :] * H_swap_mask
        # Swap the coordinate of hydrogen
        crd_1_3[..., self.hindices, :] += h_minus_value
        crd_1_3[..., self.hindices, :] += h_add_value

        new_crd = crd_1_3.reshape(crd_1_3_.shape)
        return new_crd

    def get_mass_center(self, crd_):
        """ Get the mass center of a given molecule. """
        # (B, G, A, D)
        crd = crd_.reshape((crd_.shape[0], -1, 3, crd_.shape[-1]))
        # (B, G, 1, D)
        cm = torch.sum(crd * self._mass, axis=-2) / self._mass.sum(axis=-2)
        return cm[..., None, :]

    def get_vel_force_update(self, crd0_, vel0_):
        """ Get the update of velocity and force.
        Args:
            crd0_ (Tensor): The coordinate after SETTLE in the origin axis.
            vel0_ (Tensor): The initial velocity.
        Returns:
            dv (Tensor): The constraint velocity.
            df (Tensor): The constraint force.
        """
        #pylint: disable=invalid-name
        crd0 = crd0_.reshape((crd0_.shape[0], -1, 3, 3))
        vel0 = vel0_.reshape((vel0_.shape[0], -1, 3, 3))
        # (B, C, 1)
        AB = torch.norm(crd0[..., 1, :] - crd0[..., 0, :], dim=-1, keepdim=True)
        AC = torch.norm(crd0[..., 2, :] - crd0[..., 0, :], dim=-1, keepdim=True)
        BC = torch.norm(crd0[..., 2, :] - crd0[..., 1, :], dim=-1, keepdim=True)

        # (B, C, D)
        e_AB = (crd0[..., 1, :] - crd0[..., 0, :]) / AB
        e_CA = (crd0[..., 0, :] - crd0[..., 2, :]) / AC
        e_BC = (crd0[..., 2, :] - crd0[..., 1, :]) / BC

        # (B, C)
        cosA = (AB ** 2 + AC ** 2 - BC ** 2) / 2 / AB / AC
        cosB = (AB ** 2 + BC ** 2 - AC ** 2) / 2 / AB / BC
        cosC = (BC ** 2 + AC ** 2 - AB ** 2) / 2 / BC / AC

        # (B, C)
        d = (2 * (self.mass_vertex + self.mass_base) ** 2 +
             2 * self.mass_vertex * self.mass_base * cosA * cosB * cosC -
             2 * (self.mass_base ** 2) * (cosA ** 2) -
             self.mass_vertex * (self.mass_vertex + self.mass_base) * (cosB ** 2 +cosC ** 2)
             ) * self.time_step * 0.5 / self.mass_base

        # (B, C, D)
        v_AB_ = vel0[..., 1, :] - vel0[..., 0, :]
        v_CA_ = vel0[..., 0, :] - vel0[..., 2, :]
        v_BC_ = vel0[..., 2, :] - vel0[..., 1, :]

        # (B, C, 1)
        v_AB = self.dot((e_AB, v_AB_))[..., None]
        v_CA = self.dot((e_CA, v_CA_))[..., None]
        v_BC = self.dot((e_BC, v_BC_))[..., None]

        # (B, C)
        tau_AB = (v_AB * (2 * (self.mass_vertex + self.mass_base) - self.mass_vertex * cosC ** 2) +
                  v_BC * (self.mass_base * cosC * cosA - (self.mass_vertex + self.mass_base) * cosB) +
                  v_CA * (self.mass_vertex * cosB * cosC - 2 * self.mass_base * cosA)
                  ) * self.mass_vertex / d

        tau_BC = (v_BC * ((self.mass_vertex + self.mass_base) ** 2 - (self.mass_base ** 2) * (cosA ** 2)) +
                  v_CA * self.mass_vertex * (self.mass_base * cosA * cosB -
                                             (self.mass_vertex + self.mass_base) * cosC) +
                  v_AB * self.mass_vertex * (self.mass_base * cosC * cosA -
                                             (self.mass_vertex + self.mass_base) * cosB)) / d

        tau_CA = (v_CA * (2 * (self.mass_vertex + self.mass_base) - self.mass_vertex * cosB ** 2) +
                  v_AB * (self.mass_vertex * cosB * cosC - 2 * self.mass_base * cosA) +
                  v_BC * (self.mass_base * cosA * cosB - (self.mass_vertex + self.mass_base) * cosC)
                  ) * self.mass_vertex / d

        # (B, C, D)
        df_A = tau_AB * e_AB - tau_CA * e_CA
        dv_A = self.time_step * df_A / 2 / self.mass_vertex
        df_B = tau_BC * e_BC - tau_AB * e_AB
        dv_B = self.time_step * df_B / 2 / self.mass_base
        df_C = tau_CA * e_CA - tau_BC * e_BC
        dv_C = self.time_step * df_C / 2 / self.mass_base

        # (B, A, D)
        return torch.stack((dv_A, dv_B, dv_C), dim=-2).reshape(vel0_.shape), \
            torch.stack((df_A, df_B, df_C), dim=-2).reshape(vel0_.shape)

    def forward(self,
              coordinate: Tensor,
              velocity: Tensor,
              force: Tensor,
              energy: Tensor,
              virial: Tensor = None,
              pbc_box: Tensor = None,
              step: int = 0,
              **kwargs
              ) -> Dict[str, Tensor]:
        """Constraint the water molecules.

        Args:
            coordinate (Tensor): Tensor of shape (B, A, D). Data type is float.
            velocity (Tensor): Tensor of shape (B, A, D). Data type is float.
            force (Tensor): Tensor of shape (B, A, D). Data type is float.
            energy (Tensor): Tensor of shape (B, 1). Data type is float.
            virial (Tensor): Tensor of shape (B, D). Data type is float.
            pbc_box (Tensor): Tensor of shape (B, D). Data type is float.
            step (int): Simulation step. Default: 0

        Returns:
            Dict[str, Tensor], Dictionary of output tensors.

        Symbols:
            B:  Number of walkers in simulation.
            A:  Number of atoms.
            D:  Spatial dimension of the simulation system. Usually is 3.

        """
        #pylint: disable=invalid-name
        # (B, A, D)
        crd_old_0_0 = torch.take_along_axis(self._coordinate.clone(), self.settle_index[..., None], axis=-2)
        crd_new_0_1 = torch.take_along_axis(coordinate.clone(), self.settle_index[..., None], axis=-2)
        mc = self.get_mass_center(crd_new_0_1)
        # (B, C, 4)
        transform, itransform = self.get_transform(crd_old_0_0)
        # (B, A, D)
        crd_new_0_2 = self.apply_transform(transform, crd_new_0_1)
        # (B, A, D)
        crd_new_0_3 = self._rotation(crd_old_0_0, crd_new_0_2)
        # (B, A, D)
        crd_new_0_3 = self.apply_transform(itransform, crd_new_0_3)
        # (B, A, D)
        crd_new_0_3 = self._swap_h(crd_new_0_2, crd_new_0_3)
        # (B, A, D)
        crd_new_0_3 += mc
        # (B, A, D)
        crd_new_0_3 = crd_new_0_3.reshape(crd_old_0_0.shape)

        update_crd = torch.take_along_axis(coordinate.clone(), self.settle_index[..., None], axis=-2)
        update_vel = torch.take_along_axis(velocity.clone(), self.settle_index[..., None], axis=-2)
        update_frc = torch.take_along_axis(force.clone(), self.settle_index[..., None], axis=-2)
        dv, df = self.get_vel_force_update(update_crd, update_vel)

        # (B, A, D)
        coordinate = coordinate.clone()
        velocity = velocity.clone()
        force = force.clone()

        coordinate.scatter_(-2, self.bs_index, crd_new_0_3)
        velocity.scatter_(-2, self.bs_index, update_vel + dv)
        force.scatter_(-2, self.bs_index, update_frc + df)

        return {'coordinate': coordinate,
                'velocity': velocity,
                'force': force,
                'energy': energy,
                'virial': virial,
                'pbc_box': pbc_box}

